from typing import  Dict, List
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import BaseMessage

# 内存中存储聊天历史的字典
message_stores: Dict[str, "InMemoryHistory"] = {}


class InMemoryHistory(BaseChatMessageHistory):
    """内存中的聊天历史实现"""

    messages: List[BaseMessage] = []

    def __init__(self):
        self.messages = []

    def add_messages(self, messages: List[BaseMessage]) -> None:
        """添加多条消息到存储中"""
        self.messages.extend(messages)

    def clear(self) -> None:
        """清空所有消息"""
        self.messages = []


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    """获取会话历史，如果不存在则创建新的"""
    if session_id not in message_stores:
        message_stores[session_id] = InMemoryHistory()
    return message_stores[session_id]